/*
 * admm.h
 *
 *  Created on: Jun 26, 2011
 *      Author: tqlong
 */

#ifndef ADMM_H_
#define ADMM_H_

#include <armadillo>

/**
 *  Alternate direction method of multipliers (ADMM) template
 *  solving LASSO :
 *    min   0.5 ||Ax-b||_2^2 + \lambda |x|_1,
 *  equivalently
 *    min   0.5 ||Ax-b||_2^2 + \lambda |z|_1
 *    s.t   x - z = 0
 */
class Lasso
{
public:
  typedef arma::mat mat_type;
  typedef arma::vec vec_type;
public:
  Lasso(const mat_type& A, const vec_type& b, double lambda) :
    A_(A), b_(b), lambda_(lambda)
  {
  }

  void
  search()
  {
    x_ = vec_type(x_size());
    x_.fill(0.0);
    z_ = vec_type(z_size());
    z_.fill(0.0);
    u_ = vec_type(u_size());
    u_.fill(0.0);

    eta_ = 1.0;
    tolerance_ = 1e-6;
    iteration_ = 0;

    R_ = arma::chol(trans(A_) * A_ + eta_ * arma::eye<mat_type>(x_size(),
        x_size()));
    do
      {
        proximal_op_x();

        z_prev_ = z_;
        proximal_op_z();
        dual_residual_ = z_ - z_prev_;

        residual_ = x_ - z_;
        u_ += residual_;

        iteration_++;

        std::cout << "iter=" << iteration_ << " |Ax-b|=" << arma::norm(A_ * x_
            - b_, 2) << " |x|_1=" << arma::norm(x_, 1) << " fx=" << fx()
            << std::endl;
      }
    while (iteration_ < 100 && (arma::norm(residual_, 2) > tolerance_
        || arma::norm(dual_residual_, 2) > tolerance_));
  }

  int
  x_size() const
  {
    return A_.n_cols;
  }

  int
  z_size() const
  {
    return A_.n_cols;
  }

  int
  u_size() const
  {
    return A_.n_cols;
  }

  /** x(k+1) = argmin_x 0.5||Ax-b||^2 + 0.5 eta ||x-z(k)+u(k)||^2
   *         = (A^T A + eta I)^{-1} * (A^T b + eta (z - u))
   */
  void
  proximal_op_x()
  {
    vec_type rhs = trans(A_) * b_ + eta_ * (z_ - u_);
    vec_type tmp = solve(trimatl(trans(R_)), rhs);
    x_ = solve(trimatu(R_), tmp);
  }

  void
  proximal_op_z()
  {
    vec_type v = x_ + u_;
    double shrink = lambda_ / eta_;
    for (int i = 0; i < z_size(); i++)
      {
        if (v[i] > shrink)
          z_[i] = v[i] - shrink;
        else if (v[i] < -shrink)
          z_[i] = v[i] + shrink;
        else
          z_[i] = 0;
      }
  }

  const vec_type&
  x() const
  {
    return x_;
  }

  double
  fx() const
  {
    double v = arma::norm(A_ * x_ - b_, 2);
    return 0.5 * v * v + lambda_ * arma::norm(x_, 1);
  }
protected:
  const mat_type& A_;
  const vec_type& b_;
  double lambda_, eta_, tolerance_;
  vec_type x_, z_, u_, z_prev_, residual_, dual_residual_;
  mat_type R_;
  int iteration_;
};

#endif /* ADMM_H_ */
